import os
import shutil
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QFileDialog,QMessageBox
from PyQt5.QtGui import QPixmap

"""
可视化快速分类图片
更改class_1_button等对应名称，更改类型
"""

class ImageClassifier(QWidget):
    def __init__(self):
        super().__init__()

        self.current_image_index = 0
        self.images = []
        self.image_paths = []
        self.image_label = QLabel()

        self.class_1_button = QPushButton("散货船重载")
        self.class_2_button = QPushButton("散货船轻载")
        self.class_3_button = QPushButton("集装箱船重载")
        self.class_4_button = QPushButton("集装箱船轻载")
        self.class_5_button = QPushButton("油轮重载")
        self.class_6_button = QPushButton("油轮轻载")
        self.class_7_button = QPushButton("工程船")
        self.class_8_button = QPushButton("执法船")

        self.class_1_button.clicked.connect(lambda: self.classify_image("SH_Z"))
        self.class_2_button.clicked.connect(lambda: self.classify_image("SH_Q"))
        self.class_3_button.clicked.connect(lambda: self.classify_image("JZX_Z"))
        self.class_4_button.clicked.connect(lambda: self.classify_image("JZX_Q"))
        self.class_5_button.clicked.connect(lambda: self.classify_image("YL_Z"))
        self.class_6_button.clicked.connect(lambda: self.classify_image("YL_Q"))
        self.class_7_button.clicked.connect(lambda: self.classify_image("GC"))
        self.class_8_button.clicked.connect(lambda: self.classify_image("ZF"))

        # self.previous_button = QPushButton("Previous")
        # self.next_button = QPushButton("Next")

        # self.previous_button.clicked.connect(self.previous_image)
        # self.next_button.clicked.connect(self.next_image)

        self.load_button = QPushButton("Load")
        self.load_button.clicked.connect(self.load_images)

        self.save_dir = "./ship"  # 修改默认保存路径

        layout_buttons = QHBoxLayout()
        # layout_buttons.addWidget(self.previous_button)
        # layout_buttons.addWidget(self.next_button)
        layout_buttons.addWidget(self.class_1_button)
        layout_buttons.addWidget(self.class_2_button)
        layout_buttons.addWidget(self.class_3_button)
        layout_buttons.addWidget(self.class_4_button)
        layout_buttons.addWidget(self.class_5_button)
        layout_buttons.addWidget(self.class_6_button)
        layout_buttons.addWidget(self.class_7_button)
        layout_buttons.addWidget(self.class_8_button)
        layout_buttons.addWidget(self.load_button)

        layout_main = QVBoxLayout()
        layout_main.addWidget(self.image_label)
        layout_main.addLayout(layout_buttons)

        self.setLayout(layout_main)

    def load_images(self):
        folder_path = QFileDialog.getExistingDirectory(self, "Open Folder")
        self.image_paths = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".jpg") or f.endswith(".png")]
        self.images = [QPixmap(path) for path in self.image_paths]
        self.current_image_index = 0
        self.show_image()

    def select_save_dir(self):
        self.save_dir = QFileDialog.getExistingDirectory(self, "Select Save Directory")

    def classify_image(self, class_id):
        if len(self.images) > 0 and self.save_dir is not None:
            image_path = self.image_paths[self.current_image_index]
            new_path = os.path.join(self.save_dir,class_id, os.path.basename(image_path))
            folder_path = os.path.join(self.save_dir, class_id)
            if not os.path.exists(folder_path):
                # 如果不存在则创建文件夹
                os.makedirs(folder_path)
                print(f"文件夹 '{folder_path}' 已创建")

            shutil.copyfile(image_path, new_path)
            self.next_image()
        else:
            QMessageBox.warning(self, "Warning", "Please load images and select a save directory first.")

    def next_image(self):
        if len(self.images) == self.current_image_index:
            QMessageBox.warning(self, "Warning", "打完了！！！！！！！！")

        if len(self.images) > 0:
            self.current_image_index = (self.current_image_index + 1 % len(self.images))
            self.show_image()

    def show_image(self):
        if len(self.images) > 0:
            self.image_label.setPixmap(self.images[self.current_image_index].scaledToWidth(600))

if __name__ == '__main__':
    app = QApplication([])
    window = ImageClassifier()
    window.show()
    app.exec_()